use std::any::Any;
use std::sync::Arc;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use futures::future::join_all;
use tokio::sync::mpsc;
use tokio::time::sleep;

use protoactor_rs::actor::{Actor, ActorContext, ActorError, DefaultActorContext, PID, Props, SystemMessage};
use protoactor_rs::system::ActorSystem;
use protoactor_rs::dispatcher::Dispatchers;

// 定义消息类型
#[derive(Clone)]
struct PingMessage {
    sender: PID,
    count: usize,
}

#[derive(Clone)]
struct PongMessage {
    count: usize,
}

// 定义mailbox类型
#[derive(Clone, Debug)]
enum MailboxType {
    Tokio,
    #[cfg(feature = "crossbeam_mailbox")]
    CrossbeamUnbounded,
}

impl std::fmt::Display for MailboxType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MailboxType::Tokio => write!(f, "Tokio Unbounded"),
            #[cfg(feature = "crossbeam_mailbox")]
            MailboxType::CrossbeamUnbounded => write!(f, "Crossbeam Unbounded"),
        }
    }
}

// 定义 TestCoordinatorActor
struct TestCoordinatorActor {
    batch_size: usize,
    ping_count: usize,
    actor_count: usize,
    report_tx: mpsc::Sender<TestResult>,
    mailbox_type: MailboxType,
}

impl Clone for TestCoordinatorActor {
    fn clone(&self) -> Self {
        Self {
            batch_size: self.batch_size,
            ping_count: self.ping_count,
            actor_count: self.actor_count,
            report_tx: self.report_tx.clone(),
            mailbox_type: self.mailbox_type.clone(),
        }
    }
}

struct TestResult {
    mailbox_type: String,
    actor_count: usize,
    message_count: usize,
    elapsed: Duration,
    messages_per_second: f64,
}

// 简化版本的Props创建函数
fn create_actor_props<A: Actor + Clone + 'static>(
    actor_producer: impl Fn() -> A + Send + Sync + 'static,
) -> Props {
    Props::from_producer(actor_producer)
}

#[async_trait]
impl Actor for TestCoordinatorActor {
    type Context = DefaultActorContext;

    async fn receive(&self, _ctx: &mut Self::Context, msg: Box<dyn Any + Send>) -> Result<(), ActorError> {
        if let Ok(_) = msg.downcast::<String>() {
            println!("Starting stress test with {} actors, {} pings per actor, batch size {}, mailbox type: {}",
                self.actor_count, self.ping_count, self.batch_size, self.mailbox_type);
            
            // 创建测试开始时间
            let start = Instant::now();
            
            // 获取当前系统实例
            let system = ActorSystem::new().unwrap();
            
            // 创建PingActor和PongActor
            let mut ping_actors = Vec::with_capacity(self.actor_count);
            let mut pong_actors = Vec::with_capacity(self.actor_count);
            
            // 复制值，避免闭包中捕获self引用
            let batch_size = self.batch_size;
            
            println!("Creating {} actor pairs...", self.actor_count);
            
            for i in 0..self.actor_count {
                // 创建基本Props
                let pong_props = create_actor_props(
                    || PongActor {}
                );
                
                let ping_props = create_actor_props(
                    move || PingActor { batch_size }
                );
                
                // 使用system生成actors，处理错误
                println!("Spawning actor pair {}/{}...", i+1, self.actor_count);
                let pong_pid = match system.spawn::<PongActor>(pong_props).await {
                    Ok(pid) => pid,
                    Err(e) => {
                        println!("Failed to spawn PongActor: {:?}", e);
                        return Err(ActorError::Unknown(format!("Failed to spawn PongActor: {:?}", e)));
                    },
                };
                
                let ping_pid = match system.spawn::<PingActor>(ping_props).await {
                    Ok(pid) => pid,
                    Err(e) => {
                        println!("Failed to spawn PingActor: {:?}", e);
                        return Err(ActorError::Unknown(format!("Failed to spawn PingActor: {:?}", e)));
                    },
                };
                
                pong_actors.push(pong_pid);
                ping_actors.push(ping_pid);
            }
            
            // 构建开始消息，每个PingActor会给相应的PongActor发送ping_count次消息
            let mut futures = Vec::with_capacity(self.actor_count);
            
            println!("Sending messages to all actors...");
            for i in 0..self.actor_count {
                let msg = PingMessage {
                    sender: ping_actors[i].clone(),
                    count: self.ping_count,
                };
                
                println!("Sending {} messages to actor pair {}/{}...", self.ping_count, i+1, self.actor_count);
                let future = system.request::<String>(&pong_actors[i], msg, Duration::from_secs(60));
                futures.push(future);
            }
            
            // 等待所有测试完成
            println!("Waiting for all actors to complete...");
            let results = join_all(futures).await;
            
            // 检查结果
            let mut success_count = 0;
            let mut error_count = 0;
            for (i, result) in results.iter().enumerate() {
                match result {
                    Ok(_) => {
                        success_count += 1;
                        println!("Actor pair {}/{} completed successfully", i+1, self.actor_count);
                    },
                    Err(e) => {
                        error_count += 1;
                        println!("Actor pair {}/{} failed: {:?}", i+1, self.actor_count, e);
                    },
                }
            }
            println!("Results: {} succeeded, {} failed", success_count, error_count);
            
            // 测试结束，计算结果
            let elapsed = start.elapsed();
            let total_messages = self.actor_count * self.ping_count * 2; // 每个ping都有一个pong
            let messages_per_second = total_messages as f64 / elapsed.as_secs_f64();
            
            println!("Test completed in {:?}", elapsed);
            println!("Total messages: {}", total_messages);
            println!("Messages per second: {:.2}", messages_per_second);
            
            // 报告结果
            let result = TestResult {
                mailbox_type: self.mailbox_type.to_string(),
                actor_count: self.actor_count,
                message_count: total_messages,
                elapsed,
                messages_per_second,
            };
            
            println!("Sending test result...");
            if let Err(e) = self.report_tx.send(result).await {
                println!("Failed to send test result: {}", e);
            }
            
            // 停止所有actor
            println!("Stopping all actors...");
            for (i, pid) in ping_actors.iter().chain(pong_actors.iter()).enumerate() {
                println!("Stopping actor {}/{}...", i+1, self.actor_count*2);
                let _ = system.stop(pid).await;
            }
            
            println!("Test coordinator finished");
        }
        
        Ok(())
    }
}

// 定义PingActor
#[derive(Clone)]
struct PingActor {
    batch_size: usize,
}

#[async_trait]
impl Actor for PingActor {
    type Context = DefaultActorContext;

    async fn receive(&self, _ctx: &mut Self::Context, msg: Box<dyn Any + Send>) -> Result<(), ActorError> {
        if let Ok(pong) = msg.downcast::<PongMessage>() {
            // 只需记录每100条消息，避免过多日志
            if pong.count % 100 == 0 {
                println!("PingActor received PongMessage {}", pong.count);
            }
        }
        
        Ok(())
    }
}

// 定义PongActor
#[derive(Clone)]
struct PongActor;

#[async_trait]
impl Actor for PongActor {
    type Context = DefaultActorContext;

    async fn receive(&self, ctx: &mut Self::Context, msg: Box<dyn Any + Send>) -> Result<(), ActorError> {
        if let Ok(ping) = msg.downcast::<PingMessage>() {
            // 收到ping消息，发送多个pong消息
            let sender_pid = ping.sender;
            let count = ping.count;
            
            println!("PongActor received PingMessage, will send {} responses", count);
            
            // 获取当前系统实例
            let system = ActorSystem::new().unwrap();
            
            // 分批发送pong消息，避免一次性发送太多
            let batch_size = 100;
            let batches = (count + batch_size - 1) / batch_size; // 向上取整
            
            for batch in 0..batches {
                let start = batch * batch_size;
                let end = std::cmp::min((batch + 1) * batch_size, count);
                
                println!("Sending batch {}/{} ({}-{})...", batch+1, batches, start, end-1);
                
                for i in start..end {
                    let pong = PongMessage { count: i };
                    if let Err(e) = system.send(&sender_pid, pong).await {
                        println!("Failed to send PongMessage {}: {:?}", i, e);
                    }
                }
                
                // 每批后短暂暂停，避免过于激进的发送
                if batches > 1 {
                    sleep(Duration::from_millis(10)).await;
                }
            }
            
            // 回复完成消息
            println!("PongActor finished sending all responses");
            ctx.respond("done".to_string());
        }
        
        Ok(())
    }
}

// 主函数
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("Starting stress test...");
    
    // 创建actor系统
    let system = ActorSystem::new().unwrap();
    println!("Actor system created");
    
    // 创建结果通道
    let (tx, mut rx) = mpsc::channel(100);
    println!("Result channel created");
    
    // 定义测试参数 - 使用更小的值
    let mut test_configs = Vec::new();
    
    // 添加Tokio测试配置
    // 格式: (actor_count, ping_count, batch_size, mailbox_type)
    test_configs.push((5, 100, 10, MailboxType::Tokio));
    
    // 添加Crossbeam测试配置（如果启用了crossbeam特性）
    #[cfg(feature = "crossbeam_mailbox")]
    {
        test_configs.push((5, 100, 10, MailboxType::CrossbeamUnbounded));
    }
    
    println!("Running {} test configurations", test_configs.len());
    
    // 运行测试
    for (config_idx, (actor_count, ping_count, batch_size, mailbox_type)) in test_configs.iter().enumerate() {
        println!("\nRunning test configuration {}/{}: {} actors, {} pings, batch size {}, mailbox type: {}",
            config_idx+1, test_configs.len(), actor_count, ping_count, batch_size, mailbox_type);
        
        // 创建测试协调器
        let coordinator = TestCoordinatorActor {
            batch_size: *batch_size,
            ping_count: *ping_count,
            actor_count: *actor_count,
            report_tx: tx.clone(),
            mailbox_type: mailbox_type.clone(),
        };
        
        // 创建一个新的副本用于闭包
        let coordinator_for_closure = coordinator.clone();
        let props = Props::from_producer(move || coordinator_for_closure.clone());
        
        // 生成协调器
        println!("Spawning test coordinator...");
        let pid = match system.spawn::<TestCoordinatorActor>(props).await {
            Ok(pid) => {
                println!("Test coordinator spawned successfully");
                pid
            },
            Err(e) => {
                println!("Failed to spawn TestCoordinatorActor: {:?}", e);
                let boxed_error: Box<dyn std::error::Error> = Box::new(std::io::Error::new(
                    std::io::ErrorKind::Other, 
                    format!("Failed to spawn TestCoordinatorActor: {:?}", e)
                ));
                return Err(boxed_error);
            },
        };
        
        // 发送开始消息
        println!("Sending start message to test coordinator...");
        match system.send(&pid, "start".to_string()).await {
            Ok(_) => println!("Start message sent successfully"),
            Err(e) => println!("Failed to send start message: {:?}", e),
        }
        
        // 等待测试完成
        println!("Waiting for test result...");
        if let Some(result) = rx.recv().await {
            println!("\nTest Results for {}:", result.mailbox_type);
            println!("Actor count: {}", result.actor_count);
            println!("Message count: {}", result.message_count);
            println!("Elapsed time: {:?}", result.elapsed);
            println!("Messages per second: {:.2}", result.messages_per_second);
            println!("----------------------------");
        } else {
            println!("No test result received");
        }
        
        // 停止协调器
        println!("Stopping test coordinator...");
        match system.stop(&pid).await {
            Ok(_) => println!("Test coordinator stopped successfully"),
            Err(e) => println!("Failed to stop coordinator: {:?}", e),
        }
        
        // 等待一会儿，确保所有资源都被释放
        println!("Waiting for resources to be released...");
        sleep(Duration::from_millis(500)).await;
    }
    
    println!("\nAll tests completed!");
    
    Ok(())
} 